package main

import (
	"fmt"
	"log"
	"os"
	"strings"

	"gorm.io/driver/mysql"
	"gorm.io/gen"
	"gorm.io/gorm"
)

// 数据库配置
type DBConfig struct {
	Host     string
	Port     string
	User     string
	Password string
	Database string
}

// 默认配置
var defaultConfig = DBConfig{
	Host:     "localhost",
	Port:     "3306",
	User:     "root",
	Password: "123456",
	Database: "gofilm",
}

func main() {
	// 从环境变量或命令行参数获取配置
	config := getDBConfig()

	// 连接数据库
	db := connectDB(config)

	// 创建生成器
	g := gen.NewGenerator(gen.Config{
		OutPath:           "./model/db_model", // 输出目录
		OutFile:           "query.go",         // 查询文件名
		ModelPkgPath:      "model/db_model",   // 模型包路径
		WithUnitTest:      true,               // 生成单元测试
		FieldNullable:     true,               // 字段可空
		FieldCoverable:    true,               // 字段可覆盖
		FieldSignable:     false,              // 字段可签名
		FieldWithIndexTag: true,               // 生成索引标签
		FieldWithTypeTag:  true,               // 生成类型标签
	})

	// 使用数据库连接
	g.UseDB(db)

	// 自动生成所有表的模型
	generateAllTables(g, db)

	// 自定义生成特定表（可选）
	generateCustomTables(g)

	// 执行生成
	g.Execute()

	fmt.Println("✅ GORM模型生成完成！")
	fmt.Printf("📁 输出目录: %s\n", "./model/db_model")
}

// 获取数据库配置
func getDBConfig() DBConfig {
	config := defaultConfig

	// 从环境变量读取配置
	if host := os.Getenv("DB_HOST"); host != "" {
		config.Host = host
	}
	if port := os.Getenv("DB_PORT"); port != "" {
		config.Port = port
	}
	if user := os.Getenv("DB_USER"); user != "" {
		config.User = user
	}
	if password := os.Getenv("DB_PASSWORD"); password != "" {
		config.Password = password
	}
	if database := os.Getenv("DB_NAME"); database != "" {
		config.Database = database
	}

	fmt.Printf("🔗 连接数据库: %s@%s:%s/%s\n", config.User, config.Host, config.Port, config.Database)
	return config
}

// 连接数据库
func connectDB(config DBConfig) *gorm.DB {
	dsn := fmt.Sprintf("%s:%s@tcp(%s:%s)/%s?charset=utf8mb4&parseTime=True&loc=Local",
		config.User, config.Password, config.Host, config.Port, config.Database)

	db, err := gorm.Open(mysql.Open(dsn), &gorm.Config{})
	if err != nil {
		log.Fatalf("❌ 数据库连接失败: %v", err)
	}

	fmt.Println("✅ 数据库连接成功")
	return db
}

// 自动生成所有表的模型
func generateAllTables(g *gen.Generator, db *gorm.DB) {
	// 获取所有表名
	tables := getAllTables(db)

	fmt.Printf("📋 发现 %d 个表:\n", len(tables))
	for i, table := range tables {
		fmt.Printf("  %d. %s\n", i+1, table)
	}

	// 为每个表生成模型
	for _, tableName := range tables {
		// 跳过系统表
		if isSystemTable(tableName) {
			continue
		}

		fmt.Printf("🔄 生成表模型: %s\n", tableName)

		// 生成基础模型
		model := g.GenerateModel(tableName)

		// 应用自定义配置
		model = applyCustomConfig(model, tableName)

		g.ApplyBasic(model)
	}
}

// 获取所有表名
func getAllTables(db *gorm.DB) []string {
	var tables []string

	// 查询所有表名
	rows, err := db.Raw("SHOW TABLES").Rows()
	if err != nil {
		log.Fatalf("❌ 获取表列表失败: %v", err)
	}
	defer rows.Close()

	for rows.Next() {
		var tableName string
		if err := rows.Scan(&tableName); err != nil {
			continue
		}
		tables = append(tables, tableName)
	}

	return tables
}

// 判断是否为系统表
func isSystemTable(tableName string) bool {
	systemTables := []string{
		"information_schema",
		"performance_schema",
		"mysql",
		"sys",
	}

	for _, sysTable := range systemTables {
		if strings.HasPrefix(tableName, sysTable) {
			return true
		}
	}

	return false
}

// 应用自定义配置
func applyCustomConfig(model gen.ModelOpt, tableName string) gen.ModelOpt {
	// 根据表名应用不同的配置
	switch tableName {
	case "users":
		return model.TableName("users")
	case "films":
		return model.TableName("films")
	case "film_sources":
		return model.TableName("film_sources")
	default:
		return model
	}
}

// 自定义生成特定表（高级配置）
func generateCustomTables(g *gen.Generator) {
	// 示例：为用户表添加自定义方法
	user := g.GenerateModel("users",
		gen.FieldIgnore("deleted_at"),                            // 忽略软删除字段
		gen.FieldType("status", "int"),                           // 指定字段类型
		gen.FieldTag("email", `gorm:"uniqueIndex" json:"email"`), // 自定义标签
	)

	// 添加自定义方法
	g.ApplyBasic(
		user,
		// 可以在这里添加自定义查询方法
	)

	// 示例：为影片表添加关联关系
	film := g.GenerateModel("films")
	g.ApplyBasic(film)

	// 示例：为采集源表添加配置
	filmSource := g.GenerateModel("film_sources")
	g.ApplyBasic(filmSource)
}
